#
# Created: 2025-07-27


import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import networkx as nx  # 用于网络图建模与绘制
from matplotlib import axes
from matplotlib.patches import Patch  # 用于自定义图例
from tabulate import tabulate  # 用于表格输出（美化打印）

# 设置全局字体为 Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'

# ------------------------ 核心绘图函数 ------------------------
def venn_network_plot(
        edge_data: pd.DataFrame,  # 输入边数据，必须包含 source、target、color 列
        edge_width: float = 1.0,  # 边的粗细
        edge_style: str = "line",  # 线条风格，支持 'line' 或 'curve'
        source_node_size: float = 100,  # source节点（中心）大小
        source_font_size: int = 10,  # source标签字体大小
        target_node_size: float = 5,  # target节点（外围）大小
        target_font_size: int = 10,  # target标签字体大小
        show_node_margin: bool = True,  # 是否显示节点边缘留白（线条连接位置）
        show_source_label: bool = True,  # 是否显示 source 节点标签
        show_target_label: bool = False,  # 是否显示 target 节点标签
        radius: float = 0.5,  # 控制 source 节点环形排布半径
        spring_constant: float = None,  # 节点间引力参数（影响 layout 紧凑度）
        ax: axes.Axes = None,  # 可传入已有 matplotlib 轴，否则默认新建
        label_color_map: dict = None,  # 标签颜色映射，支持每个节点单独设字体颜色
        custom_color_map: dict = None,  # 用于生成图例的 Sample 颜色映射
) -> axes.Axes:
    # 若没有传入 ax，就使用当前默认坐标轴
    if ax is None:
        ax = plt.gca()

    # 创建 NetworkX 图，读取边数据（含颜色属性）
    G = nx.from_pandas_edgelist(edge_data, edge_attr=["color"], create_using=nx.Graph())

    # 获取所有 source 节点，并安排它们均匀分布在一个圆上
    fixed_nodes = edge_data["source"].unique()
    fixed_nodes_pos = {
        node: (
            radius * np.sin(2 * np.pi * i / len(fixed_nodes)),  # x 坐标
            radius * np.cos(2 * np.pi * i / len(fixed_nodes)),  # y 坐标
        )
        for i, node in enumerate(fixed_nodes)
    }

    # 设置节点大小：source 用大圈，target 用小圈
    node_size = [
        target_node_size if node not in fixed_nodes else source_node_size
        for node in G.nodes()
    ]
    node_options = {
        "node_color": "lightgray",  # 所有节点固定为浅灰色
        "node_size": node_size,
    }

    # 设置边的颜色和样式
    edge_color = [d["color"] for _, _, d in G.edges(data=True)]
    edge_styles = {
        "line": {
            "edge_color": edge_color,
            "style": "solid",
            "alpha": 0.5,
            "width": edge_width,
        },
        "curve": {
            "edge_color": edge_color,
            "style": "solid",
            "alpha": 0.5,
            "width": edge_width,
            "arrows": True,
            "arrowsize": 10,
            "arrowstyle": None,
            "connectionstyle": "arc3,rad=0.3",  # 让线条有弯曲弧度
        }
    }

    # 用 spring_layout 算法计算节点坐标（带固定的 source 节点）
    pos = nx.spring_layout(
        G,
        pos=fixed_nodes_pos,
        fixed=fixed_nodes,
        seed=111,
        k=spring_constant,
        iterations=200,
    )

    # 绘制节点与边
    nx.draw_networkx_nodes(G, pos, **node_options)
    nx.draw_networkx_edges(G, pos, **edge_styles[edge_style])

    # 若不需要显示节点边缘，可以去除默认缩进
    if not show_node_margin:
        for patch in ax.patches:
            patch.shrinkA = patch.shrinkB = 0.0

    # 标签绘制逻辑：按颜色分组绘制，支持 source 和 target 分别控制是否显示
    if show_source_label or show_target_label:
        all_label_nodes = []
        if show_source_label:
            all_label_nodes.extend(fixed_nodes)
        if show_target_label:
            all_label_nodes.extend([n for n in G.nodes if n not in fixed_nodes])

        # 分颜色绘制，避免 label_font_color 只能统一设置的问题
        color_to_nodes = {}
        for node in all_label_nodes:
            color = label_color_map.get(node, "black") if label_color_map else "black"
            color_to_nodes.setdefault(color, []).append(node)

        for color, nodes in color_to_nodes.items():
            label_dict = {node: node for node in nodes}
            nx.draw_networkx_labels(
                G,
                pos,
                labels=label_dict,
                font_size=source_font_size,
                font_color=color,
                font_weight="bold",
                ax=ax,
            )

    # 自定义图例（每个 Sample 一个颜色块）
    if custom_color_map:
        legend_elements = [
            Patch(facecolor=color, edgecolor='none', label=label)
            for label, color in custom_color_map.items()
        ]
        ax.legend(
            handles=legend_elements,
            loc='upper right',
            bbox_to_anchor=(1.0, 1.15),  # 控制图例往上靠一些
            title="Samples",
            frameon=False
        )

    ax.set_axis_off()
    return ax


# ===================== 主程序入口 ======================
if __name__ == "__main__":
    data_file = "edge_data.csv"  # 保存的边数据文件（第一次运行请先生成）

    # 定义 source 节点（样本分类）
    Samples = ["SampleA", "SampleB", "SampleC", "SampleD", "SampleE"]

    # 自定义每个 Sample 的边颜色（浅色）
    custom_color_map = {
        "SampleA": "#b0d1a6",
        "SampleB": "#f5c5a0",
        "SampleC": "#e7a39a",
        "SampleD": "#b1a5c9",
        "SampleE": "#a0d8ef",
    }

    # 自定义每个 Sample 的标签颜色（深色）
    sample_label_color_map = {
        "SampleA": "#4f7942",
        "SampleB": "#c65f0a",
        "SampleC": "#a03530",
        "SampleD": "#5e4c8c",
        "SampleE": "#20648a",
    }

    # ------------------ 读取之前保存的 edge_data ------------------
    edge_data = pd.read_csv(data_file)
    print("[✔] 读取本地数据完成。")

    # 展示前5行边数据（美化输出）
    print(tabulate(edge_data.head(5), headers="keys", tablefmt="grid"))

    # ------------------ 调用绘图函数 ------------------
    fig, ax = plt.subplots(figsize=(6, 6))  # 创建6x6英寸图像画布
    venn_network_plot(
        edge_data,
        edge_style="curve",  # 使用曲线样式
        target_node_size=10,  # target 节点大小
        spring_constant=0.1,  # 控制布局紧凑程度
        show_source_label=True,
        show_target_label=False,
        label_color_map=sample_label_color_map,
        custom_color_map=custom_color_map,
        ax=ax
    )

    # ------------------ 保存图片 ------------------
    fig.savefig("venn_network_plot.png", dpi=300, bbox_inches='tight')
    plt.show()
